package com.hk.core.sharding;

import org.springframework.core.NamedThreadLocal;
import org.springframework.transaction.support.TransactionSynchronization;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.function.Supplier;

/**
 * @author Kevin
 * @date 2024-07-03 17:03:20
 */
public class ThreadLocalHintTokenProvider implements HintTokenProvider {

    private static final ThreadLocal<HintToken> TOKEN_THREAD_LOCAL = new NamedThreadLocal<>("ThreadLocalHintToken");

    public static void set(final HintToken hintToken) {
        TOKEN_THREAD_LOCAL.set(hintToken);
    }

    public static void clear() {
        TOKEN_THREAD_LOCAL.remove();
    }

    public static void clearWithTransaction() {
        // 如果在service类或方法加上事物注解，对于insert/update/delete语句，需要在事物执行完后，再将数据清除，不然选择的数据源会有问题
        if (TransactionSynchronizationManager.isActualTransactionActive()) {
            TransactionSynchronizationManager.registerSynchronization(new TransactionSynchronization() {

                @Override
                public void afterCompletion(int status) {
                    clear();
                }
            });
        } else {
            clear();
        }
    }

    public static <T> T invokeAndGet(HintToken token, Supplier<T> supplier) {
        try {
            set(token);
            return supplier.get();
        } finally {
            clearWithTransaction();
        }
    }

    public static void invoke(HintToken token, Runnable run) {
        try {
            set(token);
            run.run();
        } finally {
            clearWithTransaction();
        }
    }

    public static <V> Future<V> asyncInvokeAndGet(HintToken token,
                                                  ExecutorService executorService,
                                                  Supplier<V> supplier) {
        return executorService.submit(() -> {
            try {
                set(token);
                return supplier.get();
            } finally {
                clearWithTransaction();
            }
        });
    }

    public static void asyncInvoke(HintToken token, Executor executor, Runnable run) {
        executor.execute(() -> {
            try {
                set(token);
                run.run();
            } finally {
                clearWithTransaction();
            }
        });
    }

    @Override
    public HintToken get() {
        return TOKEN_THREAD_LOCAL.get();
    }
}
